import os

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments
from datasets import load_dataset
from safetensors.torch import load_file

# Check if CUDA is available and set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

os.environ["HF_DATASETS_CACHE"] = "D:/Project/AI/cache/huggingface_datasets"

# Load model and tokenizer
model_path = "D:/Project/AI/GoogleBERT/bert-base-multilingual-cased"
tokenizer = AutoTokenizer.from_pretrained(model_path)

#创建一个预训练的模型，并指定它用于 三分类 任务（num_labels=23。这就是模型的类别数，决定了模型输出层的大小（即输出三个类别的 logits）。这也是 CrossEntropyLoss 计算损失时所期望的类别数。也就是n_classes。
model = AutoModelForSequenceClassification.from_pretrained(
    model_path, num_labels=3)

# Move model to GPU if available
model = model.to(device)

# Load and preprocess data

# 使用 XNLI 数据集进行微调，XNLI 数据集是一个 三分类任务，XNLI 数据集本身是一个 跨语言的自然语言推理（NLI）任务，其中包含了 15 种语言，且每个示例都有一个标签，该标签表示句子对之间的关系，包括以下 3 个分类标签：
# 0: entailment（蕴含）
# 1: contradiction（矛盾）
# 2: neutral（中立）
dataset_path = "D:/Project/AI/datasets/facebook/xnli"
dataset = load_dataset(dataset_path, "all_languages")

# 数据预处理：tokenize
def tokenize_function(examples):
    # 提取英文翻译作为训练文本
    premises = [premise["en"] for premise in examples["premise"]]  # 获取英文的 premise
    hypotheses = [hypothesis["translation"][4] for hypothesis in examples["hypothesis"]] # 获取英文的 hypothesis（索引4是英文）
    
    # 使用 tokenizer 进行批量处理
    # 使用 tokenizer 对 premises 和 hypotheses 进行分词，padding="max_length" 表示将文本填充到最大长度，truncation=True 表示截断超过最大长度的文本。
    tokenized_output = tokenizer(
        premises, hypotheses, padding="max_length", truncation=True
    )

    # 获取标签并存储到tokenized_output
    labels = examples["label"]  # 对应每个示例的标签
    tokenized_output["labels"] = labels

    return tokenized_output

# Avoid multiprocessing (set num_proc=1)
# tokenized_dataset = dataset.map(tokenize_function, batched=True, num_proc=1)
tokenized_dataset = dataset.map(tokenize_function, batched=True)
# Move dataset to GPU (if possible) and pin memory
tokenized_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])

# Set training arguments
training_args = TrainingArguments(
    output_dir="./results", # 训练结果保存的目录。
    num_train_epochs=3, # 训练的轮数（epoch）
    per_device_train_batch_size=8, 
    per_device_eval_batch_size=8,  # 每个设备（GPU/CPU）上的批处理大小
    evaluation_strategy="epoch",  # 指定评估策略，此处为每个 epoch 后评估一次。
    logging_dir='./logs',  # 日志文件保存目录。
    logging_steps=100,  # 每 100 步记录日志。
    save_strategy="epoch",  # 每个 epoch 后保存模型。
    load_best_model_at_end=True,  # 训练结束后加载最优模型。
    metric_for_best_model="accuracy",  # 用于选择最优模型的评估指标（这里使用 accuracy）。
    fp16=True,  # Use 16-bit mixed precision training to speed up training
    dataloader_pin_memory=True,  # 将数据加载到固定内存中，加速数据加载。
    # dataloader_num_workers=0,  # Disable multiprocessing for DataLoader
)

# Train
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],  # 可选：设置测试集
)

if __name__ == "__main__":
    trainer.train()
